import time
import numpy as np
import cv2
from src.actions import SetServo, Stop, Start, TurnLeft, TurnRight, Advance
from src.scenes.base_scene import BaseScene
from src.utils import log

def extract_yellow_region(image):
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    lower_yellow = np.array([20, 100, 180])
    upper_yellow = np.array([30, 150, 255])
    mask = cv2.inRange(hsv, lower_yellow, upper_yellow)
    yellow_region = cv2.bitwise_and(image, image, mask=mask)
    gray = cv2.cvtColor(yellow_region, cv2.COLOR_BGR2GRAY)
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)

    # 强化虚线
    kernel = np.ones((7, 7), np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

    return binary

def getEdgeImg(binary_img):
    return cv2.Canny(binary_img, 50, 100)

def roi_mask(img, corner_points):
    mask = np.zeros_like(img)
    cv2.fillPoly(mask, corner_points, 255)
    return cv2.bitwise_and(img, mask)

def get_lane_lines(edges):
    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=10,
                            minLineLength=20, maxLineGap=5)
    left_lines = []
    right_lines = []

    if lines is None:
        return None, None

    for line in lines:
        x1, y1, x2, y2 = line[0]
        if x2 - x1 == 0:
            continue
        slope = (y2 - y1) / (x2 - x1)
        if slope < -0.3:
            left_lines.append(line[0])
        elif slope > 0.3:
            right_lines.append(line[0])

    return left_lines, right_lines

def average_lane_position(lines, img_width, img_height):
    if not lines:
        return None
    x_points = []
    for x1, y1, x2, y2 in lines:
        x_points += [x1, x2]
    return int(np.mean(x_points))

def get_cur_offset_x(imgOri):
    small_img = cv2.resize(imgOri, (int(imgOri.shape[1] / 4), int(imgOri.shape[0] / 4)))
    roi_binary = extract_yellow_region(small_img)
    edges = getEdgeImg(roi_binary)

    rows, cols = edges.shape
    points = np.array([[(0, rows), (0, 90), (cols, 90), (cols, rows)]])
    roi_edges = roi_mask(edges, points)

    left_lines, right_lines = get_lane_lines(roi_edges)
    img_center = cols // 2

    left_x = average_lane_position(left_lines, cols, rows)
    right_x = average_lane_position(right_lines, cols, rows)

    if left_x is not None and right_x is not None:
        lane_center = (left_x + right_x) // 2
    elif left_x is not None:
        lane_center = left_x + 100  # 估算右侧
    elif right_x is not None:
        lane_center = right_x - 100  # 估算左侧
    else:
        lane_center = img_center  # 无线，保持直行

    offset = (lane_center - img_center) * 4  # 放大因缩小了1/4
    return offset

class LF_PID(BaseScene):
    def __init__(self, memory_name, camera_info, msg_queue):
        super().__init__(memory_name, camera_info, msg_queue)
        self.forward_spd = 26

    def init_state(self):
        log.info(f'start init {self.__class__.__name__}')
        self.ctrl.execute(SetServo(servo=[90, 30]))
        return False

    def loop(self):
        ret = self.init_state()
        if ret:
            log.error(f'{self.__class__.__name__} init failed.')
            return
        frame = np.ndarray((self.height, self.width, 3), dtype=np.uint8, buffer=self.broadcaster.buf)
        log.info(f'{self.__class__.__name__} loop start')
        try:
            while True:
                if self.stop_sign.value:
                    break
                if self.pause_sign.value:
                    continue
                start = time.time()
                img_bgr = frame.copy()

                x_offset = get_cur_offset_x(img_bgr)
                log.info(f'x_offset: {x_offset}')

                if -60 < x_offset <= 60:
                    self.ctrl.execute(Advance(speed=self.forward_spd))
                elif x_offset > 60:
                    self.ctrl.execute(TurnRight(degree=min(0.45, x_offset / 300)))
                elif x_offset < -60:
                    self.ctrl.execute(TurnLeft(degree=min(0.45, -x_offset / 300)))

                log.info(f'LF cost {time.time() - start}')
        except KeyboardInterrupt:
            self.ctrl.execute(Stop())
